#!/usr/bin/env python3
# -*- coding:utf-8 -*-
u"""
2018.3.22
从MySQL表中抽取10,000行，做样本
"""
import os
from peewee import CharField
from peewee import IntegerField
from peewee import Model
from peewee import SqliteDatabase
from peewee import IntegrityError
from peewee import OperationalError

try:
    from src.db.mysqlDB import EnhancerTable as MysqlDbTable
except ImportError:
    try:
        from db.mysqlDB import EnhancerTable as MysqlDbTable
    except ImportError:
        from mysqlDB import EnhancerTable as MysqlDbTable


__author__ = "张一鸣"
__date__ = "2018.3.22"
__version__ = 0.1


__dir__ = os.path.dirname(os.path.abspath(__file__))

# Sqlite3 connector，随意指定路径
db = SqliteDatabase(os.path.join(__dir__, "FFLoop.db"))


def create_table(table):
    u"""
    检测表是否已经建立
    若没有，则新建

    并不处理任何报错
    :param table: peewee的表格模型
    :return: True 建立成功，False失败, None表格已存在
    """
    if not table.table_exists():
        table.create_table()
        return True
    else:
        return None
    return False


def drop_table(table):
    u"""
    同create_table

    并不处理任何报错
    :param table: peewee的表格模型
    :return: True 删除成功，False删除失败, None表格不存在
    """
    if table.table_exists():
        table.drop_table()
        return True
    else:
        return None
    return False


class BaseModel(Model):
    u"""
    基础数据库模型，与相应的数据库进行链接，或指定其他该数据库的参数
    其他表继承该类即可
    """
    class Meta:
        u"""
        此处设定该数据库共有的参数
        """
        database = db


class EnhancerTable(BaseModel):
    u"""
    样表
    """

    chrom = CharField(index=True)
    start = IntegerField(index=True)
    end = IntegerField(index=True)
    enhancerID = CharField(index=True)
    tissue = CharField(index=True)

    class Meta:
        u"""
        此处设定该表格参数
        """
        db_table = 'CellEnhancerTable'
        order_by = ('chrom', 'start', 'end', )
        indexes = (
            # chrom、start、end比较多的是一起查询，因此三者共同建立一个index，True表明这三个值的组合唯一，不得重复
            (('chrom', 'start', 'end'), True),
            (('chrom', 'start', 'end', 'tissue'), False)
        )


if __name__ == '__main__':
    from tqdm import trange

    def extract_data_from_mysql(number):
        u"""
        使用pyMySQL查询10000条的写法
        :param number:需要多少条
        :return: 字典构成的列表
        """
        import pymysql

        # Connect to the database
        connection = pymysql.connect(host='localhost',
                                     user='2016201050',
                                     password='2016201050',
                                     db='FFLoop',
                                     charset='utf8mb4',
                                     cursorclass=pymysql.cursors.DictCursor)

        try:
            with connection.cursor() as cursor:
                # Create a new record
                sql = "SELECT * FROM `CellEnhancerTable` LIMIT %s"
                cursor.execute(sql, (number,))
                data = cursor.fetchall()
                return data
        finally:
            connection.close()

    # 查询并且中等速度插入10000条数据
    create_table(EnhancerTable)
    with db.atomic():
        # data = extract_data_from_mysql(10000)
        for i in trange(0, 10000, 100):

            data = [x for x in MysqlDbTable.select().where(
                (MysqlDbTable.id <= i + 100) &
                (MysqlDbTable.id > i)
            ).dicts()]

            for j in data:
                try:
                    EnhancerTable.insert(**j).execute()
                except IndentationError as err:
                    print(err)
                    print(j)
                    exit(1)

    pass
